package com.gnote.langchain.config;

import com.gnote.langchain.service.MongoChatMemoryStore;
import com.gnote.langchain.service.ToolsService;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.UserMessage;
import lombok.RequiredArgsConstructor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * @author 李经纬
 */
@Configuration
@RequiredArgsConstructor
public class AiConfig {

    public interface AssistantUnique {

        String chat(@MemoryId int memoryId, @UserMessage String userMessage);

        // 流式响应
        TokenStream stream(@MemoryId int memoryId, @UserMessage String userMessage);
    }

    @Bean
    public AssistantUnique assistantUnique(ChatLanguageModel qwenChatModel,
                                           StreamingChatLanguageModel qwenStreamingChatModel,
                                           ToolsService toolsService,
                                           MongoChatMemoryStore mongoChatMemoryStore) {
        return AiServices.builder(AssistantUnique.class)
                .chatLanguageModel(qwenChatModel)
                .streamingChatLanguageModel(qwenStreamingChatModel)
                .chatMemoryProvider(memoryId ->
                        MessageWindowChatMemory.builder().maxMessages(20)
                                .chatMemoryStore(mongoChatMemoryStore)
                                .id(memoryId)
                                .build()
                )
                .tools(toolsService)
                .build();
    }
}